(* Graph representation: adjacency list as (node * neighbors list) list *)
type 'a graph = ('a * 'a list) list

(* Helper function to find neighbors of a node in the graph *)
let rec neighbors node = function
  | [] -> []
  | (n, adj) :: rest ->
      if n = node then adj else neighbors node rest

(* Function to determine the degree of a given node *)
let degree graph node =
  let nbrs = neighbors node graph in
  List.length nbrs

(* Function to generate a list of all nodes sorted by decreasing degree *)
let nodes_by_decreasing_degree graph =
  let all_nodes = List.map fst graph in
  let node_degrees = List.map (fun node -> (node, degree graph node)) all_nodes in
  (* Sort by degree descending, then by node value for stability *)
  let sorted = List.sort (fun (n1, d1) (n2, d2) ->
      if d1 <> d2 then compare d2 d1  (* descending by degree *)
      else compare n1 n2              (* ascending by node for tie-breaking *)
    ) node_degrees in
  List.map fst sorted

(* Welsh-Powell graph coloring algorithm *)
let welsh_powell_coloring graph =
  let sorted_nodes = nodes_by_decreasing_degree graph in
  
  (* Helper function to find the smallest available color for a node *)
  let rec smallest_available_color node colored_graph =
    let node_neighbors = neighbors node graph in
    let used_colors = 
      List.fold_left (fun acc neighbor ->
          try
            let color = List.assoc neighbor colored_graph in
            color :: acc
          with Not_found -> acc
        ) [] node_neighbors in
    let used_colors_set = List.sort_uniq compare used_colors in
    
    (* Find the smallest non-negative integer not in used_colors_set *)
    let rec find_smallest color =
      if List.mem color used_colors_set then
        find_smallest (color + 1)
      else
        color
    in
    find_smallest 0
  in
  
  (* Color nodes in order of decreasing degree *)
  let rec color_nodes remaining colored =
    match remaining with
    | [] -> colored
    | node :: rest ->
        let color = smallest_available_color node colored in
        color_nodes rest ((node, color) :: colored)
  in
  
  color_nodes sorted_nodes []

(* Test functions *)
let print_graph graph =
  List.iter (fun (node, neighbors) ->
      Printf.printf "%s: [%s]\n" 
        (string_of_int node)
        (String.concat "; " (List.map string_of_int neighbors))
    ) graph

let print_coloring coloring =
  let sorted_coloring = List.sort (fun (n1, _) (n2, _) -> compare n1 n2) coloring in
  List.iter (fun (node, color) ->
      Printf.printf "Node %d -> Color %d\n" node color
    ) sorted_coloring

let test_degree () =
  let g = [
    (1, [2; 3]);
    (2, [1; 3; 4]);
    (3, [1; 2; 4]);
    (4, [2; 3])
  ] in
  Printf.printf "Testing degree function:\n";
  Printf.printf "Degree of node 1: %d\n" (degree g 1);
  Printf.printf "Degree of node 2: %d\n" (degree g 2);
  Printf.printf "Degree of node 3: %d\n" (degree g 3);
  Printf.printf "Degree of node 4: %d\n" (degree g 4);
  Printf.printf "\n"

let test_nodes_by_decreasing_degree () =
  let g = [
    (1, [2; 3]);
    (2, [1; 3; 4]);
    (3, [1; 2; 4]);
    (4, [2; 3])
  ] in
  Printf.printf "Testing nodes_by_decreasing_degree:\n";
  let sorted = nodes_by_decreasing_degree g in
  Printf.printf "Sorted nodes: [%s]\n" 
    (String.concat "; " (List.map string_of_int sorted));
  Printf.printf "\n"

let test_welsh_powell () =
  (* Test case 1: Complete graph K4 *)
  let k4 = [
    (1, [2; 3; 4]);
    (2, [1; 3; 4]);
    (3, [1; 2; 4]);
    (4, [1; 2; 3])
  ] in
  
  Printf.printf "Testing Welsh-Powell on K4 (complete graph with 4 nodes):\n";
  Printf.printf "Graph:\n";
  print_graph k4;
  let coloring_k4 = welsh_powell_coloring k4 in
  Printf.printf "Coloring:\n";
  print_coloring coloring_k4;
  Printf.printf "\n";
  
  (* Test case 2: Cycle graph C5 (odd cycle) *)
  let c5 = [
    (1, [2; 5]);
    (2, [1; 3]);
    (3, [2; 4]);
    (4, [3; 5]);
    (5, [4; 1])
  ] in
  
  Printf.printf "Testing Welsh-Powell on C5 (cycle with 5 nodes):\n";
  Printf.printf "Graph:\n";
  print_graph c5;
  let coloring_c5 = welsh_powell_coloring c5 in
  Printf.printf "Coloring:\n";
  print_coloring coloring_c5;
  Printf.printf "\n";
  
  (* Test case 3: Bipartite graph *)
  let bipartite = [
    (1, [3; 4]);
    (2, [3; 4]);
    (3, [1; 2]);
    (4, [1; 2])
  ] in
  
  Printf.printf "Testing Welsh-Powell on bipartite graph:\n";
  Printf.printf "Graph:\n";
  print_graph bipartite;
  let coloring_bipartite = welsh_powell_coloring bipartite in
  Printf.printf "Coloring:\n";
  print_coloring coloring_bipartite;
  Printf.printf "\n"

(* Validation function to check if coloring is valid *)
let validate_coloring graph coloring =
  let color_of node = List.assoc node coloring in
  let rec check_edges = function
    | [] -> true
    | (node, neighbors) :: rest ->
        let node_color = color_of node in
        let neighbors_ok = List.for_all (fun neighbor ->
            try node_color <> color_of neighbor with Not_found -> true
          ) neighbors in
        neighbors_ok && check_edges rest
  in
  check_edges graph

let test_validation () =
  let g = [
    (1, [2; 3]);
    (2, [1; 3; 4]);
    (3, [1; 2; 4]);
    (4, [2; 3])
  ] in
  let coloring = welsh_powell_coloring g in
  Printf.printf "Validating coloring: %b\n" (validate_coloring g coloring)

(* Main test function *)
let () =
  test_degree ();
  test_nodes_by_decreasing_degree ();
  test_welsh_powell ();
  test_validation ();;